{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Clustered Multitask GP (w/ Pyro/GPyTorch High-Level Interface)\n",
    "\n",
    "## Introduction\n",
    "\n",
    "In this example, we use the Pyro integration for a GP model with additional latent variables.\n",
    "\n",
    "We are modelling a multitask GP in this example. Rather than assuming a linear correlation among the different tasks, we assume that there is cluster structure for the different tasks. Let's assume there are $k$ different clusters of tasks. The generative model for task $i$ is:\n",
    "\n",
    "$$\n",
    "p(\\mathbf y_i \\mid \\mathbf x_i) = \\int \\sum_{z_i=1}^k p(\\mathbf y_i \\mid \\mathbf f (\\mathbf x_i), z_i) \\: p(z_i) \\: p(\\mathbf f (\\mathbf x_i) ) \\: d \\mathbf f\n",
    "$$\n",
    "\n",
    "where $z_i$ is the cluster assignment for task $i$. There are therefore $k$ latent functions $\\mathbf f = [f_1 \\ldots f_k]$, each modelled by a GP, representing each cluster.\n",
    "\n",
    "Our goal is therefore to infer:\n",
    "\n",
    "- The latent functions $f_1 \\ldots f_k$\n",
    "- The cluster assignments $z_i$ for each task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import pyro\n",
    "import gpytorch\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# this is for running the notebook in our testing framework\n",
    "import os\n",
    "smoke_test = ('CI' in os.environ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding additional latent variables to the likelihood\n",
    "\n",
    "The standard GPyTorch variational objects will take care of inferring the latent functions $f_1 \\ldots f_k$. However, we do need to add the additional latent variables $z_i$ to the models. We will do so by creating a custom likelihood that models:\n",
    "\n",
    "$$\n",
    "\\sum_{z_i=1}^k p(\\mathbf y_i \\mid \\mathbf f (\\mathbf x_i), z_i) \\: p(z_i)\n",
    "$$\n",
    "\n",
    "GPyTorch's likelihoods are capable of modeling additional latent variables. Our custom likelihood needs to define the following three functions:\n",
    "\n",
    "- `pyro_model` (needs to call through to `super().pyro_model` at the end), which defines the prior distribution for additional latent variables\n",
    "- `pyro_guide` (needs to call through to `super().pyro_guide` at the end), which defines the variational (guide) distribution for additional latent variables\n",
    "- `forward`, which defines the observation distributions conditioned on `\\mathbf f (\\mathbf x_i)` and any additional latent variables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The pyro_model function\n",
    "\n",
    "For each task, we will model the cluster assignment with a `OneHotCategorical` variable, where each cluster has equal probability. The `pyro_model` function will make a `pyro.sample` call to this prior distribution and then call the super method:\n",
    "\n",
    "```python\n",
    "    # self.prior_cluster_logits = torch.zeros(num_tasks, num_clusters)\n",
    "\n",
    "    def pyro_model(self, function_dist, target):\n",
    "        cluster_assignment_samples = pyro.sample(\n",
    "            self.name_prefix + \".cluster_logits\",  # self.name_prefix is added by PyroGP\n",
    "            pyro.distributions.OneHotCategorical(logits=self.prior_cluster_logits).to_event(1)\n",
    "        )\n",
    "        return super().pyro_model(\n",
    "            function_dist,\n",
    "            target,\n",
    "            cluster_assignment_samples=cluster_assignment_samples\n",
    "        )\n",
    "```\n",
    "\n",
    "Note that we are adding an additional argument `cluster_assignment_samples` to the `super().pyro_model` call. This will pass the cluster assignment samples to the `forward` call, which is necessary for inference.\n",
    "\n",
    "### The pyro_guide function\n",
    "\n",
    "For each task, the variational (guide) diustribution will also be a `OneHotCategorical` variable, which will be defined by the parameter `self.variational_cluster_logits`. The `pyro_guide` function will make a `pyro.sample` call to this prior distribution and then call the super method:\n",
    "\n",
    "```python\n",
    "    def pyro_guide(self, function_dist, target):\n",
    "        pyro.sample(\n",
    "            self.name_prefix + \".cluster_logits\",  # self.name_prefix is added by PyroGP\n",
    "            pyro.distributions.OneHotCategorical(logits=self.variational_cluster_logits).to_event(1)\n",
    "        )\n",
    "        return super().pyro_guide(function_dist, target)\n",
    "```\n",
    "\n",
    "Note that we are adding an additional argument `cluster_assignment_samples` to the `super().pyro_model` call. This will pass the cluster assignment samples to the `forward` call, which is necessary for inference.\n",
    "\n",
    "\n",
    "### The forward function\n",
    "\n",
    "The `pyro_model` fuction passes the additional keyword argument `cluster_assignment_samples` to the `forward` call. Therefore, our forward method will define the conditional probability $p(\\mathbf y_i \\mid \\mathbf f(\\mathbf x), z_i)$, where $\\mathbf f(\\mathbf x)$ corresponds to the variable `function_samples` and $z_i$ corresponds to the variable `cluster_assignment_samples`.\n",
    "\n",
    "In our example $p(\\mathbf y_i \\mid \\mathbf f(\\mathbf x), z_i)$ corresponds to a Gaussian noise model.\n",
    "\n",
    "```python\n",
    "    # self.raw_noise is the Gaussian noise parameter\n",
    "    # function_samples is `n x k`\n",
    "    # cluster_assignment_samples is `k x t`, where `t` is the number of tasks\n",
    "\n",
    "    def forward(self, function_samples, cluster_assignment_samples):\n",
    "        return pyro.distributions.Normal(\n",
    "            loc=(function_samples.unsqueeze(-2) * cluster_assignment_samples).sum(-1),\n",
    "            scale=torch.nn.functional.softplus(self.raw_noise).sqrt()\n",
    "        ).to_event(1)\n",
    "        # The to_event call is necessary because we are returning a multitask distribution,\n",
    "        # where each task dimension corresponds to each of the `t` tasks\n",
    "```\n",
    "\n",
    "This is all we need for inference! However, if we want to use this model to make predictions, the `cluster_assignment_samples` keyword argument will not be passed into the function. Therefore, we need to make sure that `forward` can handle both inference and predictions:\n",
    "\n",
    "\n",
    "```python\n",
    "    def forward(self, function_samples, cluster_assignment_samples=None):\n",
    "        if cluster_assignment_samples is None:\n",
    "            # We'll get here at prediction time\n",
    "            # We'll use the variational distribution when making predictions\n",
    "            cluster_assignment_samples = pyro.sample(\n",
    "                self.name_prefix + \".cluster_logits\", self._cluster_dist(self.variational_cluster_logits)\n",
    "            )\n",
    "            \n",
    "        return pyro.distributions.Normal(\n",
    "            loc=(function_samples.unsqueeze(-2) * cluster_assignment_samples).sum(-1),\n",
    "            scale=torch.nn.functional.softplus(self.raw_noise).sqrt()\n",
    "        ).to_event(1)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClusterGaussianLikelihood(gpytorch.likelihoods.Likelihood):\n",
    "    def __init__(self, num_tasks, num_clusters):\n",
    "        super().__init__()\n",
    "        \n",
    "        # These are parameters/buffers for the cluster assignment latent variables\n",
    "        self.register_buffer(\"prior_cluster_logits\", torch.zeros(num_tasks, num_clusters))\n",
    "        self.register_parameter(\"variational_cluster_logits\", torch.nn.Parameter(torch.randn(num_tasks, num_clusters)))\n",
    "        \n",
    "        # The Gaussian observational noise\n",
    "        self.register_parameter(\"raw_noise\", torch.nn.Parameter(torch.tensor(0.0)))\n",
    "        \n",
    "        # Other info\n",
    "        self.num_tasks = num_tasks\n",
    "        self.num_clusters = num_clusters\n",
    "        self.max_plate_nesting = 1\n",
    "\n",
    "    def pyro_guide(self, function_dist, target):\n",
    "        # Here we add the extra variational distribution for the cluster latent variable\n",
    "        pyro.sample(\n",
    "            self.name_prefix + \".cluster_logits\",  # self.name_prefix is added by PyroGP\n",
    "            pyro.distributions.OneHotCategorical(logits=self.variational_cluster_logits).to_event(1)\n",
    "        )\n",
    "        return super().pyro_guide(function_dist, target)\n",
    "\n",
    "    def pyro_model(self, function_dist, target):\n",
    "        # Here we add the extra prior distribution for the cluster latent variable\n",
    "        cluster_assignment_samples = pyro.sample(\n",
    "            self.name_prefix + \".cluster_logits\",  # self.name_prefix is added by PyroGP\n",
    "            pyro.distributions.OneHotCategorical(logits=self.prior_cluster_logits).to_event(1)\n",
    "        )\n",
    "        return super().pyro_model(function_dist, target, cluster_assignment_samples=cluster_assignment_samples)\n",
    "\n",
    "    def forward(self, function_samples, cluster_assignment_samples=None):\n",
    "        # For inference, cluster_assignment_samples will be passed in\n",
    "        # This bit of code is for when we use the likelihood in the predictive mode\n",
    "        if cluster_assignment_samples is None:\n",
    "            cluster_assignment_samples = pyro.sample(\n",
    "                self.name_prefix + \".cluster_logits\", self._cluster_dist(self.variational_cluster_logits)\n",
    "            )\n",
    "            \n",
    "        # Now we return the observational distribution, based on the function_samples and cluster_assignment_samples\n",
    "        res = pyro.distributions.Normal(\n",
    "            loc=(function_samples.unsqueeze(-2) * cluster_assignment_samples).sum(-1),\n",
    "            scale=torch.nn.functional.softplus(self.raw_noise).sqrt()\n",
    "        ).to_event(1)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constructing the PyroGP model\n",
    "\n",
    "The PyroGP model is essentially the same as the model we used in the simple example, except for two changes\n",
    "\n",
    "- We now will use our more complicated `ClusterGaussianLikelihood`\n",
    "- The latent function should be vector valued to correspond to the `k` latent functions. As a result, we will learn a batched variational distribution, and use a `MultitaskVariationalStrategy` to convert the batched variational distribution into a `MultitaskMultivariateNormal` distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClusterMultitaskGPModel(gpytorch.models.pyro.PyroGP):\n",
    "    def __init__(self, train_x, train_y, num_functions=2, reparam=False):\n",
    "        num_data = train_y.size(-2)\n",
    "\n",
    "        # Define all the variational stuff\n",
    "        inducing_points = torch.linspace(0, 1, 64).unsqueeze(-1)\n",
    "        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(\n",
    "            num_inducing_points=inducing_points.size(-2),\n",
    "            batch_shape=torch.Size([num_functions])\n",
    "        )\n",
    "        \n",
    "        # Here we're using a MultitaskVariationalStrategy - so that the output of the\n",
    "        # GP latent function is a MultitaskMultivariateNormal\n",
    "        variational_strategy = gpytorch.variational.MultitaskVariationalStrategy(\n",
    "            gpytorch.variational.VariationalStrategy(self, inducing_points, variational_distribution),\n",
    "            num_tasks=num_functions,\n",
    "        )\n",
    "\n",
    "        # Standard initializtation\n",
    "        likelihood = ClusterGaussianLikelihood(train_y.size(-1), num_functions)\n",
    "        super().__init__(variational_strategy, likelihood, num_data=num_data, name_prefix=str(time.time()))\n",
    "        self.likelihood = likelihood\n",
    "        self.num_functions = num_functions\n",
    "\n",
    "        # Mean, covar\n",
    "        self.mean_module = gpytorch.means.ZeroMean()\n",
    "        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean_x = self.mean_module(x)\n",
    "        covar_x = self.covar_module(x)\n",
    "        res = gpytorch.distributions.MultivariateNormal(mean_x, covar_x)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This model can now be used to perform inference on cluster assignments, as well as make predictions using the inferred cluster assignments!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
